# -*- coding: utf-8 -*-
"""
Created on Mon Oct  8 14:32:52 2018

@author: luolei

训练模型
"""
import torch
import json
import os
import sys

sys.path.append('../..')

from lib import *
from lib.modeling import *


def save_models(nn_model, train_loss_record, verify_loss_record):
	"""保存模型文件"""
	
	# 保存模型文件
	torch.save(nn_model.state_dict(), '../../file/model/nn_state_dict_{}_{}.pth'.format(target_cols, selected_cols))
	
	# 保存模型结构参数
	model_struc_params = {
		'nn_model': {
			'input_size': nn_model.input_size,
			'hidden_sizes': nn_model.hidden_sizes,
			'output_size': nn_model.output_size
		}
	}
	
	with open('../../file/model/nn_struc_params_{}_{}.json'.format(target_cols, selected_cols), 'w') as f:
		json.dump(model_struc_params, f)
	
	# 损失函数记录
	train_loss_list = [float(p.detach().cpu().numpy()) for p in train_loss_record]
	verify_loss_list = [float(p.cpu().numpy()) for p in verify_loss_record]
	
	with open('../../file/model/nn_train_loss.json', 'w') as f:
		json.dump(train_loss_list, f)
	with open('../../file/model/nn_verify_loss.json', 'w') as f:
		json.dump(verify_loss_list, f)
	
	
if __name__ == '__main__':
	#%% 是否使用GPU
	use_cuda = torch.cuda.is_available()
	
	#%% 载入数据集，构建训练和验证集样本
	trainloader, verifyloader, X_train, y_train, X_verify, y_verify, total_cols_n = build_train_and_verify_datasets()
	
	#%% 构造神经网络模型
	input_size = X_train.shape[1]
	output_size = y_train.shape[1]
	hidden_sizes = [2 * input_size, 2 * input_size, input_size, 2 * output_size, 2 * output_size, 2 * output_size]
	nn_model = NN(input_size, hidden_sizes, output_size)

	if use_cuda:
		torch.cuda.empty_cache()
		trainloader = [(train_x.cuda(), train_y.cuda()) for (train_x, train_y) in trainloader]
		verifyloader = [(verify_x.cuda(), verify_y.cuda()) for (verify_x, verify_y) in verifyloader]
		nn_model = nn_model.cuda()

	#%% 指定优化器
	optimizer = torch.optim.Adam(
		nn_model.parameters(),
		lr = lr
	)

	#%% 模型训练和保存
	train_loss_record, verify_loss_record = [], []

	for epoch in range(epochs):
		# 训练集
		nn_model.train()
		for train_x, train_y in trainloader:
			y_train_p = nn_model(train_x)
			y_train_t = train_y
			train_loss_fn = criterion(y_train_p, y_train_t)

			optimizer.zero_grad()
			train_loss_fn.backward()
			optimizer.step()

		train_loss_record.append(train_loss_fn)

		# 验证集
		nn_model.eval()
		with torch.no_grad():
			for verify_x, verify_y in verifyloader:
				y_verify_p = nn_model(verify_x)
				y_verify_t = verify_y
				verify_loss_fn = criterion(y_verify_p, y_verify_t)
			verify_loss_record.append(verify_loss_fn)

		if (epoch + 1) % 50 == 0:
			print(epoch, train_loss_fn, verify_loss_fn)

		# 保存模型
		if (epoch + 1) % 500 == 0:
			save_models(nn_model, train_loss_record, verify_loss_record)

	save_models(nn_model, train_loss_record, verify_loss_record)
	
	
	
	
	
	
	
	
	
	


